# The three-Body Problem

from vpython import * 

# Some important constant
G = 6.7e-11                  # Gravitational constant
Msun = 2e30                  # The mass of the Sun
Rsun = 7e8                   # The radius of the Sun
L = 1.5e11                   # The average distance from the Sun to Earth 
Vsun = sqrt(G*Msun/Rsun)/20  # The first universe Speed

# Create canvas and Camera perspective
scene.width = 1120
scene.height = 700

scene.camera.pos = vector(1.5*L, 0.3*L, 1.5*L)      # set the position of the camera
scene.camera.axis = vector(-1.5*L, -0.3*L, -1.5*L) # set the direction of the camera

# plot three plane and xyz axis
wallx = box(pos = vector(0, 0, 0), size = vector(2*L, 10000, 2*L), color = color.gray(0.5), opacity = 0.4, shininess = 0)
wally = box(pos = vector(0, 0, 0), size = vector(2*L, 2*L, 10000), color = color.gray(0.5), opacity = 0.4, shininess = 0)
wallz = box(pos = vector(0, 0, 0), size = vector(10000, 2*L, 2*L), color = color.gray(0.5), opacity = 0.4, shininess = 0)

xaxis = curve(color = color.gray(1), radius = 0.3*Rsun)
xaxis.append(vec(0, 0, 0))
xaxis.append(vec(L, 0, 0))
yaxis = curve(color = color.gray(1), radius = 0.3*Rsun)
yaxis.append(vec(0, 0, 0))
yaxis.append(vec(0, L, 0))
zaxis = curve(color = color.gray(1), radius = 0.3*Rsun)
zaxis.append(vec(0, 0, 0))
zaxis.append(vec(0, 0, L))

Xlabel = label(text = "+X", height = 20, opacity = 0, box = False)   # define X label
Ylabel = label(text = "+Y", height = 20, opacity = 0, box = False)   # define Y label
Zlabel = label(text = "+Z", height = 20, opacity = 0, box = False)   # define Z label

Xlabel.pos = vec(1.05*L, 0, 0)                                # define X label position
Ylabel.pos = vec(0, 1.05*L, 0)                                # define Y label position
Zlabel.pos = vec(0, 0, 1.05*L)                                # define Z label position

# draw three stars
body1 = sphere(pos = vector(-L/2, 0, 0), radius = 10*Rsun, color = color.blue, make_trail = True, trail_type = 'points', interval = 5, retain = 500)
body2 = sphere(pos = vector(L/4, 0, 0), radius = 7*Rsun, color = color.yellow, make_trail = True, trail_type = 'points', interval = 5, retain = 500)
body3 = sphere(pos = vector(L/2, 0, 0), radius = 5*Rsun, color = color.red, make_trail = True, trail_type = 'points', interval = 5, retain = 500)

# ################################## interaction part
begin = False

def Restart(b):
    global begin
    begin = not begin
    if begin:
        b.text = 'RESTART'
    else:
        b.text = 'START'
        
b1 = button(text = 'start', pos = scene.title_anchor, bind = Restart)

def sigmaF(ma, mb, mc, Rac, Rbc):
    F1 = G*ma*mc*Rac.hat/mag(Rac)**2
    F2 = G*mb*mc*Rbc.hat/mag(Rbc)**2
    return F1 + F2

body1.m = 1*Msun
body2.m = 0.7*Msun
body3.m = 0.5*Msun

body1.v = vector(0.8*Vsun, 0, 0.8*Vsun)
body2.v = vector(-0.8*Vsun, 0.8*Vsun, 0)
body3.v = vector(0, 0.9*Vsun, 0.9*Vsun)

body1.a = vector(0, 0, 0)
body2.a = vector(0, 0, 0)
body3.a = vector(0, 0, 0)

t = 0
dt = 1e3 
while True:
    rate(1000)
    
    # for the body1
    body1.a = sigmaF(body2.m, body3.m, body1.m, body2.pos - body1.pos, body3.pos - body1.pos) / body1.m
    body1.v = body1.v + body1.a*dt
    body1.pos = body1.pos + body1.v*dt
    
    # for the body2
    body2.a = sigmaF(body1.m, body3.m, body2.m, body1.pos - body2.pos, body3.pos - body2.pos) / body2.m
    body2.v = body2.v + body2.a*dt
    body2.pos = body2.pos + body2.v*dt
    
    # for the body3
    body3.a = sigmaF(body1.m, body2.m, body3.m, body1.pos - body3.pos, body2.pos - body3.pos) / body3.m
    body3.v = body3.v + body3.a*dt
    body3.pos = body3.pos + body3.v*dt
    
    t = t + dt